{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import tqdm\n",
    "from tqdm import tqdm_notebook, notebook\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import rankeval\n",
    "import scipy as sc\n",
    "import time\n",
    "import hyperopt\n",
    "import copy\n",
    "import random as rng\n",
    "import lightgbm as lgb\n",
    "import cython\n",
    "\n",
    "import seaborn as sns\n",
    "\n",
    "from sklearn.metrics import classification_report, confusion_matrix\n",
    "from sklearn.utils import class_weight\n",
    "\n",
    "from rankeval.dataset import Dataset\n",
    "from rankeval.model import RTEnsemble\n",
    "from rankeval.metrics import Metric, NDCG, DCG, MAP\n",
    "from rankeval.metrics import NDCG\n",
    "import rankeval.analysis.statistical\n",
    "import rankeval.metrics\n",
    "from rankeval.analysis.feature import feature_importance\n",
    "\n",
    "from hyperopt import base\n",
    "import gc\n",
    "base.have_bson = False\n",
    "\n",
    "%load_ext Cython\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from collections import defaultdict\n",
    "\n",
    "class SDCG(Metric):\n",
    "\n",
    "    def __init__(self, name='SDCG', cutoff=None, implementation=\"flat\"):\n",
    "\n",
    "        super(SDCG, self).__init__(name)\n",
    "        self.cutoff = cutoff\n",
    "        self.implementation = implementation\n",
    "\n",
    "    def eval(self, dataset, y_pred):\n",
    "\n",
    "        return super(SDCG, self).eval(dataset, y_pred)\n",
    "\n",
    "    def eval_per_query(self, y, y_pred):\n",
    "        idx_y_pred_sorted = np.argsort(y_pred,kind='stable')[::-1]\n",
    "        if self.cutoff is not None:\n",
    "            idx_y_pred_sorted = idx_y_pred_sorted[:self.cutoff]\n",
    "\n",
    "        discount = np.log2(np.arange(2, idx_y_pred_sorted.size + 2))\n",
    "\n",
    "        if self.implementation == \"flat\":\n",
    "            gain = y[idx_y_pred_sorted]\n",
    "        elif self.implementation == \"exp\":\n",
    "            gain = np.exp2(y[idx_y_pred_sorted]) - 1.0\n",
    "\n",
    "        dcg = (gain / discount).sum()\n",
    "        return dcg\n",
    "\n",
    "    def __str__(self):\n",
    "        s = self.name\n",
    "        if self.cutoff is not None:\n",
    "            s += \"@{}\".format(self.cutoff)\n",
    "        return s\n",
    "\n",
    "class SNDCG(Metric):\n",
    "\n",
    "    def __init__(self, name='SNDCG', cutoff=None, no_relevant_results=1.0,\n",
    "                 implementation=\"exp\"):\n",
    "\n",
    "        super(self.__class__, self).__init__(name)\n",
    "        self.cutoff = cutoff\n",
    "        self.no_relevant_results = no_relevant_results\n",
    "        self.implementation = implementation\n",
    "        self.dcg = SDCG(cutoff=self.cutoff,\n",
    "                       implementation=self.implementation)\n",
    "\n",
    "        self._current_dataset = None\n",
    "        self._current_rel_qid = None\n",
    "        self._cache_idcg_score = defaultdict(int)\n",
    "\n",
    "    def eval(self, dataset, y_pred):\n",
    "        # used to cache ideal DCG scores on a dataset basis\n",
    "        self._current_dataset = dataset\n",
    "        self._current_rel_qid = 0\n",
    "\n",
    "        # Compute the ideal DCG scores only once and cache them\n",
    "        if self._current_dataset not in self._cache_idcg_score:\n",
    "\n",
    "            idcg_score = np.ndarray(shape=dataset.n_queries, dtype=np.float32)\n",
    "            for rel_id, (qid, q_y, _) in enumerate(\n",
    "                    self.query_iterator(dataset, dataset.y)):\n",
    "                idcg_score[rel_id] = self.dcg.eval_per_query(q_y, q_y)\n",
    "\n",
    "            self._cache_idcg_score[self._current_dataset] = idcg_score\n",
    "\n",
    "        return super(self.__class__, self).eval(dataset, y_pred)\n",
    "\n",
    "    def eval_per_query(self, y, y_pred):\n",
    "        dcg_score = self.dcg.eval_per_query(y, y_pred)\n",
    "\n",
    "        if self._current_rel_qid is not None:\n",
    "            idcg_score = \\\n",
    "                self._cache_idcg_score[self._current_dataset][self._current_rel_qid]\n",
    "            self._current_rel_qid += 1\n",
    "        else:\n",
    "            idcg_score = self.dcg.eval_per_query(y, y)\n",
    "\n",
    "        if idcg_score != 0:\n",
    "            ndcg = dcg_score / idcg_score\n",
    "        else:\n",
    "            ndcg = self.no_relevant_results\n",
    "        return ndcg\n",
    "\n",
    "    def __str__(self):\n",
    "        s = self.name\n",
    "        if self.cutoff is not None:\n",
    "            s += \"@{}\".format(self.cutoff)\n",
    "        return s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "istella_name = \"ISTELLA-FULL\"\n",
    "msn_name = \"MSN30KF1\"\n",
    "istella_path = './data/ISTELLA/'\n",
    "msn_path = './data/MSLR/Fold1'\n",
    "\n",
    "datasets_name = [msn_name, istella_name]\n",
    "datasets_pairs = [(msn_name, msn_path), (istella_name, istella_path)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rankeval_datasets = {}\n",
    "for dataset_name, dataset_path in datasets_pairs:\n",
    "    rankeval_datasets[dataset_name] = dict([\n",
    "        (split, Dataset.load(os.path.join(dataset_path, split + '.txt'), name=\"{} {}\".format(dataset_name, split.capitalize()))) for split in [ 'vali', 'test']\n",
    "])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rankeval_models = {}\n",
    "\n",
    "for dataset_name in datasets_name:\n",
    "    path_model = './data/rankers/{}-lgb_model.txt'.format(dataset_name)\n",
    "    rankeval_models[dataset_name] = RTEnsemble(path_model, format='QuickRank')\n",
    "    print('Model trained on {} has {} trees'.format(dataset_name, rankeval_models[dataset_name].n_trees))\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "feature_importances = {}\n",
    "for dataset_name in tqdm.notebook.tqdm(datasets_name):\n",
    "    f_importances = feature_importance(rankeval_models[dataset_name], rankeval_datasets[dataset_name]['vali'], metric=None, normalize=True)\n",
    "    importances = f_importances[0].sortby(f_importances[0], ascending=False)\n",
    "    features_by_importance = importances.coords['feature'].values\n",
    "    feature_importances[dataset_name] = features_by_importance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_pred_rankeval = {}\n",
    "y_partial_pred_rankeval = {}\n",
    "y_pred_along_training = {}\n",
    "for dataset_name in tqdm.notebook.tqdm(datasets_name):   \n",
    "    y_pred_rankeval[dataset_name] = {}\n",
    "    y_partial_pred_rankeval[dataset_name] = {}\n",
    "    y_pred_along_training[dataset_name] = {}\n",
    "\n",
    "    for split in tqdm.notebook.tqdm(rankeval_datasets[dataset_name], leave=False):\n",
    "        y_pred_rankeval[dataset_name][split], y_partial_pred_rankeval[dataset_name][split], _ = rankeval_models[dataset_name].score(rankeval_datasets[dataset_name][split], detailed=True)\n",
    "        y_pred_along_training[dataset_name][split] = y_partial_pred_rankeval[dataset_name][split].cumsum(axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "classification_train = {}\n",
    "classification_vali = {}\n",
    "classification_test = {}\n",
    "\n",
    "y_pred_rankeval_train = {}\n",
    "y_pred_along_train = {}\n",
    "y_pred_rankeval_test = {}\n",
    "y_pred_along_test = {}\n",
    "y_pred_rankeval_vali = {}\n",
    "y_pred_along_vali = {}\n",
    "for dataset_name in datasets_name:\n",
    "    mask_validation = np.random.choice(a=[True, False, False, False], size=(rankeval_datasets[dataset_name]['test'].n_queries,))\n",
    "    mask_test = ~ mask_validation\n",
    "\n",
    "    mask_validation_docs = np.zeros((rankeval_datasets[dataset_name]['test'].n_instances,), dtype=np.bool)\n",
    "    for rel_id, (qid, start, end) in enumerate(rankeval_datasets[dataset_name]['test'].query_iterator()):\n",
    "        if mask_validation[rel_id]:\n",
    "            mask_validation_docs[start:end] = True\n",
    "    mask_test_docs = ~ mask_validation_docs\n",
    "\n",
    "    validation_qids = rankeval_datasets[dataset_name]['test'].query_ids[mask_validation]\n",
    "    test_qids = rankeval_datasets[dataset_name]['test'].query_ids[mask_test]\n",
    "\n",
    "    classification_train[dataset_name] = copy.deepcopy(rankeval_datasets[dataset_name]['vali'])\n",
    "    classification_vali[dataset_name] = rankeval_datasets[dataset_name]['test'].subset(validation_qids)\n",
    "    classification_test[dataset_name] = rankeval_datasets[dataset_name]['test'].subset(test_qids)\n",
    "\n",
    "    y_pred_rankeval_train[dataset_name] = y_pred_rankeval[dataset_name]['vali']\n",
    "    y_pred_along_train[dataset_name] = y_pred_along_training[dataset_name]['vali']\n",
    "    \n",
    "    y_pred_rankeval_test[dataset_name] = y_pred_rankeval[dataset_name]['test'][mask_test_docs]\n",
    "    y_pred_along_test[dataset_name] = y_pred_along_training[dataset_name]['test'][mask_test_docs, :]\n",
    "    \n",
    "    y_pred_rankeval_vali[dataset_name] = y_pred_rankeval[dataset_name]['test'][mask_validation_docs]\n",
    "    y_pred_along_vali[dataset_name] = y_pred_along_training[dataset_name]['test'][mask_validation_docs, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "del rankeval_datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(gc.collect())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_lgb_parameters(p):\n",
    "    # parameters to cast\n",
    "    for p_name in [\"min_data_in_leaf\", \"num_leaves\", \"max_position\", \"eval_at\"]:\n",
    "        if p_name in p:\n",
    "            p[p_name] = int(p[p_name])\n",
    "\n",
    "    params = {\n",
    "        'boosting_type': 'gbdt',\n",
    "        'objective': 'binary',\n",
    "        'max_position': 10,\n",
    "        'metric': 'binary_logloss',\n",
    "    }\n",
    "    params.update(**p)\n",
    "\n",
    "    other = {}\n",
    "    if \"num_boost_round\" in params:\n",
    "        other[\"num_boost_round\"] = params.pop(\"num_boost_round\")\n",
    "    else:\n",
    "        other[\"num_boost_round\"] = 1000\n",
    "\n",
    "    if \"early_stopping_rounds\" in params:\n",
    "        other[\"early_stopping_rounds\"] = params.pop(\"early_stopping_rounds\")\n",
    "    else:\n",
    "        other[\"early_stopping_rounds\"] = 100\n",
    "\n",
    "    return params, other"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_lgb_objective(lgb_datasets, metric, debug=None):\n",
    "\n",
    "    def lgb_objective(p):\n",
    "        try:\n",
    "            params, other = get_lgb_parameters(p)\n",
    "\n",
    "            evals_result = {}\n",
    "            lgb_model = lgb.train(\n",
    "                params, \n",
    "                lgb_datasets['train'],\n",
    "                valid_sets=[lgb_datasets['vali']],\n",
    "                valid_names=['vali'],\n",
    "                evals_result=evals_result,\n",
    "                verbose_eval=False,\n",
    "                **other\n",
    "            )\n",
    "\n",
    "            \n",
    "            loss = evals_result['vali'][metric][lgb_model.best_iteration-1]\n",
    "            \n",
    "            # debug\n",
    "            if debug is not None:\n",
    "                if \"valid_loss\" not in debug or debug[\"valid_loss\"] is None or debug[\"valid_loss\"] > loss:\n",
    "                    debug[\"p\"] = p\n",
    "                    debug[\"params\"] = params\n",
    "                    debug[\"other\"] = other\n",
    "                    debug[\"valid_loss\"] = loss\n",
    "\n",
    "                    debug[\"model\"] = lgb_model\n",
    "                    debug[\"num_trees\"] = lgb_model.best_iteration\n",
    "\n",
    "            return {\n",
    "                \"loss\": loss,\n",
    "                \"status\": hyperopt.STATUS_OK\n",
    "            }\n",
    "\n",
    "        except Exception as e:\n",
    "            print(e)\n",
    "            return {\n",
    "                \"status\": hyperopt.STATUS_FAIL,\n",
    "                \"exception\": str(e)\n",
    "            }\n",
    "\n",
    "    return lgb_objective"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_params_space(num_boost_round=1500, num_leaves=[64], max_depth=[8] ,early_stopping_rounds=100):\n",
    "    \n",
    "    # parameters space\n",
    "    params_space = {\n",
    "        'learning_rate': hyperopt.pyll.base.scope.maximum(\n",
    "            hyperopt.hp.quniform(\"learning_rate\", 10, 100, 5) / 100,\n",
    "            0.1\n",
    "        ),\n",
    "        'num_leaves': hyperopt.hp.choice('num_leaves', num_leaves),\n",
    "        'max_depth': hyperopt.hp.choice('max_depth', max_depth),\n",
    "        \n",
    "        'min_data_in_leaf': hyperopt.hp.quniform('min_data_in_leaf', 10, 10000, 10),\n",
    "        \n",
    "        'min_sum_hessian_in_leaf': hyperopt.pyll.base.scope.maximum(\n",
    "            hyperopt.hp.quniform('min_sum_hessian_in_leaf', 1, 1e5, 100) / 1000,\n",
    "            0.001\n",
    "        ),\n",
    "        \n",
    "        'lambda_l1': hyperopt.pyll.base.scope.maximum(\n",
    "            hyperopt.hp.quniform('lambda_l1', 0, 100, 2) / 10,\n",
    "            0.1\n",
    "        ),\n",
    "\n",
    "        'lambda_l2': hyperopt.pyll.base.scope.maximum(\n",
    "            hyperopt.hp.quniform('lambda_l2', 0, 100, 2) / 10,\n",
    "            0.0\n",
    "        ),\n",
    "        \n",
    "        'min_gain_to_split': hyperopt.pyll.base.scope.maximum(\n",
    "            hyperopt.hp.quniform('min_gain_to_split', 0, 100, 2) / 10,\n",
    "            0.0\n",
    "        ),\n",
    "        \n",
    "        'boosting_type': 'gbdt',\n",
    "        'early_stopping_rounds': early_stopping_rounds,\n",
    "        'num_boost_round': num_boost_round,\n",
    "        'feature_pre_filter': False,\n",
    "        'verbosity' : -1\n",
    "    }\n",
    "    \n",
    "    return params_space"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_stats(y_true, y_pred):\n",
    "    print(classification_report(y_true, y_pred))\n",
    "    cm = confusion_matrix(y_true, y_pred)\n",
    "    ax= plt.subplot()\n",
    "    sns.heatmap(cm, annot=True, ax = ax, cmap='Greens', fmt='g') #annot=True to annotate cells\n",
    "\n",
    "    # labels, title and ticks\n",
    "    ax.set_xlabel('Predicted labels')\n",
    "    ax.set_ylabel('True labels')\n",
    "    ax.set_title('Confusion Matrix')\n",
    "    ax.xaxis.set_ticklabels(['exit', 'continue']) \n",
    "    ax.yaxis.set_ticklabels(['exit', 'continue'])\n",
    "    plt.show()\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dcg_10 = SDCG(cutoff=10, implementation='exp')\n",
    "ndcg_10 = SNDCG(cutoff=10, implementation='exp', no_relevant_results=1.0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_features(dataset, scores, sentinel_tree, base_features ,rank_features=1, feat_importances=None):\n",
    "    sentinel_pos = sentinel_tree + 1\n",
    "    extra=2\n",
    "    \n",
    "    if feat_importances is not None:\n",
    "        data_features = base_features\n",
    "    \n",
    "    else:\n",
    "        data_features = dataset.n_features\n",
    "    \n",
    "    score_features = rank_features #sentinel_pos \n",
    "    n_feats = data_features + score_features  + extra + 1\n",
    "    n_inst = dataset.n_instances\n",
    "    rank_idx = np.zeros((n_inst, score_features))\n",
    "    features = np.zeros(shape=(n_inst, n_feats))\n",
    "    \n",
    "    if feat_importances is not None:\n",
    "        features[:, :data_features] = dataset.X[:, feat_importances[:data_features]]\n",
    "    else:\n",
    "        features[:, :data_features] = dataset.X[:]\n",
    "    \n",
    "    for tree in range(sentinel_pos - score_features ,sentinel_pos):\n",
    "        for rel_id, (qid, start, end) in enumerate(dataset.query_iterator()):\n",
    "            query_scores = scores[start:end, tree]\n",
    "            query_range = query_scores.max() - query_scores.min()\n",
    "            norm_query_scores = (query_scores - query_scores.min()) / (query_range) if query_range != 0 else query_scores - query_scores.min()\n",
    "            query_ranks = np.argsort(query_scores, kind='stable')[::-1]\n",
    "            col = tree - (sentinel_pos - score_features)\n",
    "            for rank, doc in enumerate(query_ranks):\n",
    "                rank_idx[start + doc, col] = rank\n",
    "            \n",
    "            features[start:end, -1] = float(end - start)\n",
    "            features[start:end, data_features+rank_features + 0] = query_scores # extra 1\n",
    "            features[start:end, data_features+rank_features + 1] = norm_query_scores # extra 2\n",
    "    \n",
    "    features[:, data_features:data_features + rank_features] = rank_idx\n",
    "    return features\n",
    "\n",
    "def get_oracle(dataset, scores, k):\n",
    "    mask = np.zeros(shape=(dataset.n_instances,))\n",
    "    for rel_id, (qid, start, end) in enumerate(dataset.query_iterator()):\n",
    "        final_ranks = np.argsort(scores[start:end, -1], kind='stable')[::-1]\n",
    "        true_rel = dataset.y[start : end]\n",
    "        top_k = final_ranks[:k]\n",
    "        top_k = np.array([r for r in top_k if true_rel[r] > 0 ], dtype=int)\n",
    "        mask[start:end][top_k] = True \n",
    "    \n",
    "    return mask\n",
    "\n",
    "def class_weight_by_query(dataset, binary_labels):\n",
    "    weights = np.zeros((dataset.n_instances))\n",
    "    for rel_id, (qid, start, end) in enumerate(dataset.query_iterator()):\n",
    "        query_labels = binary_labels[start : end]\n",
    "        query_size = end - start\n",
    "        n_pos = query_labels.sum()\n",
    "        n_neg = query_size - n_pos\n",
    "        counts = np.array([n_neg, n_pos])\n",
    "        w  = np.divide(query_size, 2*counts, out=np.zeros((2,)), where=counts!=0)\n",
    "        for doc in range(start, end):\n",
    "            label = binary_labels[doc].astype(int)\n",
    "            weights[doc] = w[label]\n",
    "        \n",
    "        \n",
    "    return weights\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#best features: msn: 54 (136), istella: 118 (220)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "k = 10\n",
    "window = 0\n",
    "n_regressors = 10\n",
    "\n",
    "models = {}\n",
    "params = {}\n",
    "valid_losses = {}\n",
    "for dataset_name in tqdm.notebook.tqdm(datasets_name):\n",
    "    base_features = [f for f in range(10, classification_train[dataset_name].n_features, classification_train[dataset_name].n_features // 6)]\n",
    "    base_features.append(classification_train[dataset_name].n_features)\n",
    "    margins = [5]\n",
    "    base_features = [54, 118]\n",
    "    \n",
    "    models[dataset_name] = {}\n",
    "    params[dataset_name] = {}\n",
    "    valid_losses[dataset_name] = {}\n",
    "    NUM_TRIALS = 100\n",
    "    POSITIONS = [50,100,200]\n",
    "\n",
    "    for ex in margins:\n",
    "        margin = k + ex\n",
    "        models[dataset_name][margin] = {}\n",
    "        params[dataset_name][margin] = {}\n",
    "        valid_losses[dataset_name][margin] = {}\n",
    "        for features in base_features:\n",
    "            models[dataset_name][margin][features] = {}\n",
    "            params[dataset_name][margin][features] = {}\n",
    "            valid_losses[dataset_name][margin][features] = {}\n",
    "            for position in POSITIONS:\n",
    "                tree = position - 1\n",
    "\n",
    "                X_train = get_features(classification_train[dataset_name], \n",
    "                                       y_pred_along_train[dataset_name], \n",
    "                                       sentinel_tree=tree, \n",
    "                                       base_features=features, \n",
    "                                       feat_importances=feature_importances[dataset_name])\n",
    "\n",
    "                X_test =  get_features(classification_test[dataset_name],  \n",
    "                                       y_pred_along_test[dataset_name],  \n",
    "                                       sentinel_tree=tree, \n",
    "                                       base_features=features, \n",
    "                                       feat_importances=feature_importances[dataset_name])\n",
    "\n",
    "                X_vali =  get_features(classification_vali[dataset_name],  \n",
    "                                       y_pred_along_vali[dataset_name],  \n",
    "                                       sentinel_tree=tree, \n",
    "                                       base_features=features, \n",
    "                                       feat_importances=feature_importances[dataset_name])\n",
    "\n",
    "\n",
    "                y_train = get_oracle(classification_train[dataset_name], y_pred_along_train[dataset_name], margin)\n",
    "                y_test =  get_oracle(classification_test[dataset_name],  y_pred_along_test[dataset_name],  margin)\n",
    "                y_vali =  get_oracle(classification_vali[dataset_name],  y_pred_along_vali[dataset_name],  margin)\n",
    "\n",
    "\n",
    "                w_train = np.exp2(classification_train[dataset_name].y) * (y_train)\n",
    "                w_test  = np.exp2(classification_test[dataset_name].y)  * (y_test)\n",
    "                w_vali  = np.exp2(classification_vali[dataset_name].y)  * (y_vali)\n",
    "\n",
    "                w_train[w_train == 0] = 1\n",
    "                w_test[w_test == 0]   = 1\n",
    "                w_vali[w_vali == 0]   = 1\n",
    "\n",
    "                w_train = w_train  * class_weight_by_query(classification_train[dataset_name], y_train)\n",
    "                w_test  = w_test   * class_weight_by_query(classification_test[dataset_name], y_test)   \n",
    "                w_vali  = w_vali   * class_weight_by_query(classification_vali[dataset_name], y_vali)  \n",
    "\n",
    "\n",
    "\n",
    "\n",
    "                data = [('train', X_train , y_train, w_train), ('vali', X_vali, y_vali, w_vali), ('test', X_test, y_test, w_test)]\n",
    "\n",
    "\n",
    "\n",
    "                lgb_datasets = dict([\n",
    "                    (dataset_name, lgb.Dataset(data=X, label=y, weight= w, silent=False)) \n",
    "                for dataset_name, X, y, w in data\n",
    "                ])\n",
    "\n",
    "                del data\n",
    "\n",
    "                opt_res = {}\n",
    "                results = {}\n",
    "                print(margin, features, position)\n",
    "                for num_trees in [n_regressors]:\n",
    "                    params_space = get_params_space(num_boost_round=num_trees, num_leaves=[64], max_depth=[8], early_stopping_rounds=10)\n",
    "                    # trials object to track progress\n",
    "                    bayes_trials = hyperopt.Trials()\n",
    "                    opt_res[num_trees] = {}\n",
    "                    # optimize\n",
    "                    np.random.seed(random_seed)\n",
    "                    best = hyperopt.fmin(\n",
    "                        fn=get_lgb_objective(lgb_datasets, metric='binary_logloss', debug=opt_res[num_trees]),\n",
    "                        space=params_space,\n",
    "                        algo=hyperopt.tpe.suggest,\n",
    "                        max_evals=NUM_TRIALS,\n",
    "                        trials=bayes_trials,\n",
    "                        rstate=np.random.RandomState(random_seed),\n",
    "                    )\n",
    "                    results[num_trees] = {\n",
    "                        'best_params': hyperopt.space_eval(params_space, best),\n",
    "                    }\n",
    "                models[dataset_name][margin][features][position] = opt_res[num_trees]['model']\n",
    "                params[dataset_name][margin][features][position] = opt_res[num_trees]['p']\n",
    "                valid_losses[dataset_name][margin][features][position] = opt_res[num_trees]['valid_loss']\n",
    "                del lgb_datasets\n",
    "                gc.collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "speedups = {}\n",
    "deltas = {}\n",
    "sizes = {}\n",
    "ndcgs = {}\n",
    "steps = [v for v in range(10, 80, 10)]\n",
    "for dataset_name in datasets_name:\n",
    "    print(dataset_name)\n",
    "    dcg_unpruned = dcg_10.eval(classification_test[dataset_name], y_pred_rankeval_test[dataset_name])[1]\n",
    "    idcg = dcg_10.eval(classification_test[dataset_name], classification_test[dataset_name].y)[1]\n",
    "    ndcg_unpruned = dcg_unpruned / idcg\n",
    "    ndcg_unpruned[np.isnan(ndcg_unpruned)] = 1.0\n",
    "    unpruned_mean = ndcg_unpruned.mean()\n",
    "    base_traversals = classification_test[dataset_name].n_instances * rankeval_models[dataset_name].n_trees\n",
    "    print('ndcg unpruned = {:.4f}\\n'.format(ndcg_unpruned.mean()))\n",
    "    speedups[dataset_name] = {}\n",
    "    deltas[dataset_name] = {}\n",
    "    sizes[dataset_name] = {}\n",
    "    ndcgs[dataset_name] = {}\n",
    "    csizes = np.zeros((classification_test[dataset_name].n_queries,))\n",
    "    for ex in tqdm.notebook.tqdm(margins, leave=False):\n",
    "        margin = k + ex\n",
    "        speedups[dataset_name][margin] = {}\n",
    "        deltas[dataset_name][margin] = {}\n",
    "        sizes[dataset_name][margin] = {}\n",
    "        ndcgs[dataset_name][margin] = {}\n",
    "        for features in tqdm.notebook.tqdm(base_features, leave=False):\n",
    "            speedups[dataset_name][margin][features]= {}\n",
    "            deltas[dataset_name][margin][features] = {}\n",
    "            sizes[dataset_name][margin][features] = {}\n",
    "            ndcgs[dataset_name][margin][features] = {}\n",
    "            for position in tqdm.notebook.tqdm(POSITIONS, leave=False):\n",
    "                sentinel_tree = position - 1\n",
    "                speedups[dataset_name][margin][features][position] = []\n",
    "                deltas[dataset_name][margin][features][position] = []\n",
    "                sizes[dataset_name][margin][features][position] = []\n",
    "                ndcgs[dataset_name][margin][features][position] = {}\n",
    "                for threshold in tqdm.notebook.tqdm(steps, leave=False):\n",
    "                    t = threshold / 100\n",
    "                    bst = models[dataset_name][margin][features][position]\n",
    "\n",
    "                    X_test =  get_features(classification_test[dataset_name],  \n",
    "                                           y_pred_along_test[dataset_name],  \n",
    "                                           sentinel_tree=sentinel_tree, \n",
    "                                           base_features=features,\n",
    "                                           feat_importances=feature_importances[dataset_name])\n",
    "\n",
    "                    y_test =  get_oracle(classification_test[dataset_name],  y_pred_along_test[dataset_name],  margin)\n",
    "\n",
    "                    y_pred = bst.predict(X_test) > t\n",
    "\n",
    "                    mask_pruning = ~ y_pred\n",
    "                    y_pred_rankeval_pruned = copy.deepcopy(y_pred_rankeval_test[dataset_name])\n",
    "                    y_pred_rankeval_pruned[mask_pruning] += (sentinel_tree - 10000)\n",
    "                    dcg_pruned = dcg_10.eval(classification_test[dataset_name], y_pred_rankeval_pruned)[1]\n",
    "                    ndcg_pruned = dcg_pruned / idcg\n",
    "                    ndcg_pruned[np.isnan(ndcg_pruned)] = 1.0\n",
    "                    pruned_mean = ndcg_pruned.mean()\n",
    "\n",
    "                    delta = - (unpruned_mean - pruned_mean) / unpruned_mean\n",
    "                    done_traversals = classification_test[dataset_name].n_instances * n_regressors + (classification_test[dataset_name].n_instances * position) + (y_pred.sum() * ( rankeval_models[dataset_name].n_trees - position))\n",
    "                    speedup = base_traversals / done_traversals \n",
    "\n",
    "                    for rel_id, (qid, start, end) in enumerate(classification_test[dataset_name].query_iterator()):\n",
    "                        y_pred_query = y_pred[start : end]\n",
    "                        ndocs = y_pred_query.sum()\n",
    "                        csizes[rel_id] = ndocs\n",
    "\n",
    "                    ndcgs[dataset_name][margin][features][position][threshold] = ndcg_pruned    \n",
    "                    deltas[dataset_name][margin][features][position].append(delta * 100)\n",
    "                    speedups[dataset_name][margin][features][position].append(speedup)\n",
    "                    sizes[dataset_name][margin][features][position].append(np.percentile(csizes, 50))\n",
    "\n",
    "                    #print('margin = {}'.format(margin))\n",
    "                    #print('position = {}'.format(position))\n",
    "                    #print('features = {}'.format(features))\n",
    "                    #print('confidence = {:.2f}'.format(t))\n",
    "                    #print('ndcg pruned =  {:.4f}'.format(ndcg_pruned.mean()))\n",
    "                    print('delta ndcg  = {:.4f} %'.format(delta * 100))\n",
    "                    print('speedup = {:.2f} x'.format(speedup))\n",
    "                    \n",
    "                    #print_stats(y_test, y_pred)\n",
    "                    \n",
    "                    print('-----------------------------------------------------')\n",
    "\n",
    "\n",
    "gc.collect()\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ept_speedups = {}\n",
    "ept_deltas = {}\n",
    "ept_sizes = {}\n",
    "ept_ndcg = {}\n",
    "ept_steps = [v for v in range(10, 120, 10)]\n",
    "for dataset_name in datasets_name:\n",
    "    print(dataset_name)\n",
    "    dcg_unpruned = dcg_10.eval(classification_test[dataset_name], y_pred_rankeval_test[dataset_name])[1]\n",
    "    idcg = dcg_10.eval(classification_test[dataset_name], classification_test[dataset_name].y)[1]\n",
    "    ndcg_unpruned = dcg_unpruned / idcg\n",
    "    ndcg_unpruned[np.isnan(ndcg_unpruned)] = 1.0\n",
    "    unpruned_mean = ndcg_unpruned.mean()\n",
    "    base_traversals = classification_test[dataset_name].n_instances *  rankeval_models[dataset_name].n_trees\n",
    "    print('ndcg unpruned = {:.4f}\\n'.format(ndcg_unpruned.mean()))\n",
    "    \n",
    "    ept_speedups[dataset_name] = {}\n",
    "    ept_deltas[dataset_name] = {}\n",
    "    ept_sizes[dataset_name] = {}\n",
    "    ept_ndcg[dataset_name] = {}\n",
    "    \n",
    "    for position in POSITIONS:          \n",
    "        ept_speedups[dataset_name][position] = []\n",
    "        ept_deltas[dataset_name][position] = []\n",
    "        ept_sizes[dataset_name][position] = []\n",
    "        ept_ndcg[dataset_name][position] = {}\n",
    "        for threshold in ept_steps:\n",
    "            y_pred = np.zeros(shape=(classification_test[dataset_name].n_instances,))\n",
    "            for rel_id, (qid, start, end) in enumerate(classification_test[dataset_name].query_iterator()):\n",
    "                query_scores = y_pred_along_test[dataset_name][start:end, position - 1]\n",
    "                final_ranks = np.argsort(query_scores, kind='stable')[::-1]\n",
    "                top_k = final_ranks[:15]\n",
    "                score_kth = query_scores[top_k[-1]]\n",
    "                t = score_kth - (threshold / 100)\n",
    "                y_pred[start:end] = query_scores >= t\n",
    "\n",
    "            mask_pruning = ~ (y_pred.astype(bool))\n",
    "            y_pred_rankeval_pruned = copy.deepcopy(y_pred_rankeval_test[dataset_name])\n",
    "            y_pred_rankeval_pruned[mask_pruning] += (sentinel_tree - 10000)\n",
    "            dcg_pruned = dcg_10.eval(classification_test[dataset_name], y_pred_rankeval_pruned)[1]\n",
    "            ndcg_pruned = dcg_pruned / idcg\n",
    "            ndcg_pruned[np.isnan(ndcg_pruned)] = 1.0\n",
    "            pruned_mean = ndcg_pruned.mean()\n",
    "\n",
    "            delta = - (unpruned_mean - pruned_mean) / unpruned_mean\n",
    "            done_traversals = (classification_test[dataset_name].n_instances * position) + (y_pred.sum() * (rankeval_models[dataset_name].n_trees - position))\n",
    "            speedup = base_traversals  / done_traversals        \n",
    "\n",
    "            for rel_id, (qid, start, end) in enumerate(classification_test[dataset_name].query_iterator()):\n",
    "                y_pred_query = y_pred[start : end]\n",
    "                ndocs = y_pred_query.sum()\n",
    "                csizes[rel_id] = ndocs\n",
    "\n",
    "            ept_ndcg[dataset_name][position][threshold] = ndcg_pruned\n",
    "\n",
    "\n",
    "            ept_deltas[dataset_name][position].append( delta * 100)\n",
    "            ept_speedups[dataset_name][position].append(speedup)\n",
    "            ept_sizes[dataset_name][position].append(csizes.mean())\n",
    "            print(position, threshold/100, pruned_mean, delta*100 ,speedup)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams.update({'font.size': 18})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "losses = {}\n",
    "for dataset_name in datasets_name:\n",
    "    losses[dataset_name] = {}\n",
    "    losses[dataset_name][50] =  [valid_losses[dataset_name][15][f][50] for f in base_features]\n",
    "    losses[dataset_name][100] = [valid_losses[dataset_name][15][f][100] for f in base_features]\n",
    "    losses[dataset_name][200] = [valid_losses[dataset_name][15][f][200] for f in base_features]\n",
    "    print(base_features)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for dataset_name in datasets_name:\n",
    "    fig, ax1 = plt.subplots(figsize=(6,4))\n",
    "    for position in POSITIONS:\n",
    "        ax1.plot(base_features, losses[dataset_name][position], label=position, marker='.')\n",
    "    plt.grid()\n",
    "    plt.title(dataset_name)\n",
    "    plt.legend()\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_n_features = {msn_name: 54, istella_name: 118}\n",
    "best_margin = 15\n",
    "position_colors = {50: 'blue', 100: 'red', 200:'green'}\n",
    "models_color = {'EPT': 'red', 'LEAR': 'blue'}\n",
    "from sklearn import metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## calssificatore performances\n",
    "# ROC curve\n",
    "for dataset_name in datasets_name:\n",
    "    fig, ax1 = plt.subplots(figsize=(6,4))\n",
    "    for position in POSITIONS:\n",
    "        bst = models[dataset_name][best_margin][best_n_features[dataset_name]][position]\n",
    "        X =  get_features(classification_test[dataset_name], \n",
    "                          y_pred_along_test[dataset_name],\n",
    "                          sentinel_tree=position-1, \n",
    "                          base_features=best_n_features[dataset_name],\n",
    "                          feat_importances=feature_importances[dataset_name])\n",
    "        y_true =  get_oracle(classification_test[dataset_name],  y_pred_along_test[dataset_name],  best_margin)\n",
    "        y_pred = bst.predict(X)\n",
    "        fpr, tpr, threshold = metrics.roc_curve(y_true, y_pred)\n",
    "        #plt.title('Receiver Operating Characteristic')\n",
    "        ax1.plot(fpr, tpr, label = position, color=position_colors[position])\n",
    "\n",
    "    ax1.plot([0, 1], [0, 1],'k--')\n",
    "    ax1.legend(loc = 'lower right')\n",
    "    ax1.set_xlim([0, 1])\n",
    "    ax1.set_ylim([0, 1])\n",
    "    ax1.set_ylabel('True Positive Rate')\n",
    "    ax1.set_xlabel('False Positive Rate')\n",
    "    ax1.grid(ls=(0, (5, 5)), alpha=0.5)\n",
    "    plt.tight_layout()\n",
    "    plt.title(dataset_name)\n",
    "    plt.show()\n",
    "    #fig.savefig('../images/{}_ROC_curve.pdf'.format(dataset_name), format='pdf', bbox_inches='tight')\n",
    "\n",
    "    fig, ax1 = plt.subplots(figsize=(6,4))\n",
    "    # P-R curve\n",
    "    for position in POSITIONS:\n",
    "        bst = models[dataset_name][best_margin][best_n_features[dataset_name]][position]\n",
    "        X =  get_features(classification_test[dataset_name], \n",
    "                          y_pred_along_test[dataset_name],\n",
    "                          sentinel_tree=position-1, \n",
    "                          base_features=best_n_features[dataset_name],\n",
    "                          feat_importances=feature_importances[dataset_name])\n",
    "        y_true =  get_oracle(classification_test[dataset_name],  y_pred_along_test[dataset_name],  best_margin)\n",
    "        y_pred = bst.predict(X)\n",
    "        precision, recall, _ = metrics.precision_recall_curve(y_true, y_pred)\n",
    "        base = y_true[y_true >= 0.99].size / y_true.size\n",
    "        ax1.plot([0, 1], [base, base],'k--')\n",
    "        ax1.plot(recall, precision, marker=',', label=position, color=position_colors[position])\n",
    "        # axis labels\n",
    "    ax1.set_xlabel('Recall')\n",
    "    ax1.set_ylabel('Precision')\n",
    "    # show the legend\n",
    "    ax1.legend()\n",
    "    ax1.grid(ls=(0, (5, 5)), alpha=0.5)\n",
    "    plt.tight_layout()\n",
    "    plt.title(dataset_name)\n",
    "    plt.show()\n",
    "    #fig.savefig('../images/{}_PR_curve.pdf'.format(dataset_name), format='pdf', bbox_inches='tight')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for dataset_name in datasets_name:\n",
    "    ## EPT performances\n",
    "    fig, ax1 = plt.subplots(figsize=(6,4))\n",
    "    for position in POSITIONS:\n",
    "        ept_delta = ept_deltas[dataset_name][position]\n",
    "        ept_speedup = ept_speedups[dataset_name][position]\n",
    "        ax1.plot(ept_speedup, ept_delta, color=position_colors[position], marker='.', label='s = '+str(position),markersize=18)\n",
    "    ax1.set_xlabel(r'speedup')\n",
    "    #ax1.set_ylabel(r'$\\Delta$ NDCG@10 (%)')\n",
    "    ax1.grid(ls=(0, (5, 5)), alpha=0.5)\n",
    "    plt.legend(loc=3)\n",
    "    #plt.ylim((-0.5, 0.05))\n",
    "    #plt.xlim((1.4, 3.750))\n",
    "    plt.tight_layout()\n",
    "    plt.title(dataset_name)\n",
    "    plt.show()\n",
    "    #fig.savefig('../images/{}_EPT_tradeoff.pdf'.format(dataset_name), format='pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## LEAR performances\n",
    "for dataset_name in datasets_name:\n",
    "    fig, ax1 = plt.subplots(figsize=(6,4))\n",
    "    for position in POSITIONS:\n",
    "        ept_delta = ept_deltas[dataset_name][position]\n",
    "        ept_speedup = ept_speedups[dataset_name][position]\n",
    "        ax1.plot(speedups[dataset_name][best_margin][best_n_features[dataset_name]][position], deltas[dataset_name][best_margin][best_n_features[dataset_name]][position], color=position_colors[position], marker='.', label='s = '+str(position), markersize=18)\n",
    "    ax1.set_xlabel(r'speedup')\n",
    "    ax1.set_ylabel(r'$\\Delta$ NDCG@10 (%)')\n",
    "    ax1.grid(ls=(0, (5, 5)), alpha=0.5)\n",
    "    plt.legend(loc=3)\n",
    "    plt.tight_layout()\n",
    "    #plt.ylim((-0.5, 0.05))\n",
    "    #plt.xlim((2.1, 4.1))\n",
    "    plt.title(dataset_name)\n",
    "    plt.show()\n",
    "#fig.savefig('../images/{}_LEAR_tradeoff.pdf'.format(dataset_name), format='pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "best_position_ept = {msn_name: 200, istella_name: 200}\n",
    "best_position_lear = {msn_name: 50, istella_name: 100}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Best vs best\n",
    "for dataset_name in datasets_name:\n",
    "    fig, ax1 = plt.subplots(figsize=(6,4))\n",
    "    ax1.plot(speedups[dataset_name][best_margin][best_n_features[dataset_name]][best_position_lear[dataset_name]], deltas[dataset_name][best_margin][best_n_features[dataset_name]][best_position_lear[dataset_name]], color=models_color['LEAR'], marker='.', label='LEAR(s = {})'.format(best_position_lear[dataset_name]), markersize=18)\n",
    "    ax1.plot(ept_speedups[dataset_name][best_position_ept[dataset_name]], ept_deltas[dataset_name][best_position_ept[dataset_name]], color=models_color['EPT'], marker='.', label='EPT(s = {})'.format(best_position_ept[dataset_name]), markersize=18)\n",
    "    ax1.set_xlabel(r'speedup')\n",
    "    ax1.set_ylabel(r'$\\Delta$ NDCG@10 (%)')\n",
    "    ax1.grid(ls=(0, (5, 5)), alpha=0.5)\n",
    "    plt.legend(loc=3)\n",
    "    plt.tight_layout()\n",
    "    #plt.ylim((-0.30, 0.05))\n",
    "    #plt.xlim((1.25, 4.2))\n",
    "    plt.title(dataset_name)\n",
    "    plt.show()\n",
    "    #fig.savefig('../images/{}_comparison.pdf'.format(dataset_name), format='pdf', bbox_inches='tight')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
